import openai
openai.api_type = "azure"
openai.api_base = "https://chatgpt-modelops.openai.azure.com/"
openai.api_version = "2023-03-15-preview"
openai.api_key = "e53f8d329e3045e7922c6b105e3b01c0"


SYSTEM_PROMPT = "You are an excellent professional information extraction system. " \
                "The task is to locate and extract named entities offrom text and categorized their relations" \
                "The output should be in the following format:" \
                "{'entity':{'entity label1':[e1,e2,...],'entity label2':[e3,e4,...],...}," \
                "'relation':{'relation label1':[(e1,e2),(e3,e4),...],'relation label2':[(e5,e6),...]}}" \
                "where the relation label1:(e1,e2) means that there is relation1 between head entity e1 and tail entity e2" \

USER_PROMPT_1 = "Are you clear about your role?"

ASSISTANT_PROMPT_1 = "Sure, I'm ready to help you with your information extraction task. " \
                     "Please provide me with the necessary information to get started."

GUIDELINES_PROMPT = (
    "Entity labels: {},\n"
    "Relation labels: {}.\n"
    "Input: {}.\n"
    "Output:\n"
)

GUIDELINES_PROMPT_EXAMPLES = (
    "Entity labels: {},\n"
    "Relation labels: {}.\n"
    "Examples: {}.\n"
    "Input: {}.\n"
    "Output:\n"
)


def openai_chat_completion_response(final_prompt):
  response = openai.ChatCompletion.create(
      engine="modelops-chatgpt-deploy",
      messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": USER_PROMPT_1},
                {"role": "assistant", "content": ASSISTANT_PROMPT_1},
                {"role": "user", "content": final_prompt}
      ],
      temperature=0.7,
      max_tokens=800,
      top_p=0.95,
      frequency_penalty=0,
      presence_penalty=0,
      stop=None
            )

  return response['choices'][0]['message']['content'].strip(" \n")

entity_labels=["Task","Method","Metric","Material","OtherScientificTerm","Generic"]
relation_labels=["USED-FOR","FEATURE-OF","HYPONYM-OF","PART-OF","EVALUATE-FOR","COMPARE","CONJUNCTION"]
my_sentence='This paper presents an algorithm for computing optical flow , shape , motion , lighting , and albedo ' \
            'from an image sequence of a rigidly-moving Lambertian object under distant illumination '
examples=[]
example="'Input': 'The problem is formulated in a manner that subsumes structure from motion , multi-view stereo , " \
        "and photo-metric stereo as special cases .'\n" \
        "'Output':{'Entity': {'Generic':['problem'], 'Material':['motion', 'multi-view stereo', 'photo-metric stereo']}," \
        " 'Relation': {'CONJUNCTION':[('motion', 'multi-view stereo'),('multi-view stereo', 'motion')," \
        "('multi-view stereo','photo-metric stereo'),('photo-metric stereo', 'multi-view stereo')]}"
examples.append(example)
if examples:
    final_prompt=GUIDELINES_PROMPT_EXAMPLES.format(','.join(entity_labels),','.join(relation_labels),'\n'.join(examples),my_sentence)
else:
    final_prompt = GUIDELINES_PROMPT.format(','.join(entity_labels),','.join(relation_labels),my_sentence)
result = openai_chat_completion_response(final_prompt)
print(result)